import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.pqc.jcajce.spec.MLKEMParameterSpec;
import org.bouncycastle.pqc.jcajce.spec.*;
import org.bouncycastle.pqc.jcajce.*;

import javax.crypto.*;
import java.security.*;

public class JcaMLKEMExample {

    public static void main(String[] args) throws Exception {
        // 1. 注册 Provider
        Security.addProvider(new BouncyCastleProvider());
        Security.addProvider(new BouncyCastlePQCProvider());

        // 2. 生成 ML-KEM-512 密钥对
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("ML-KEM-512", "BCPQC");
        KeyPair keyPair = kpg.generateKeyPair();
        PublicKey pk = keyPair.getPublic();
        PrivateKey sk = keyPair.getPrivate();

        // 3. 封装（Encapsulate）—— 由对方操作
        KeyGenerator kGen = KeyGenerator.getInstance("ML-KEM-512", "BCPQC");
        kGen.init(MLKEMParameterSpec.ml_kem_512, pk, null);   // pk 为目标公钥
        SecretKeyWithEncapsulation secEnc = (SecretKeyWithEncapsulation) kGen.generateKey();
        byte[] cipherText = secEnc.getEncapsulation();      // 发送给私钥持有者
        byte[] sharedKeyA = secEnc.getEncoded();            // 共享密钥（K）

        // 4. 解封装（Decapsulate）—— 由私钥持有者操作
        Cipher kemCipher = Cipher.getInstance("ML-KEM-512", "BCPQC");
        kemCipher.init(Cipher.DECRYPT_MODE, sk);
        byte[] sharedKeyB = kemCipher.doFinal(cipherText);  // 得到相同的 K

        // 5. 验证两端共享密钥一致
        System.out.println("Alice 端共享密钥: " + bytesToHex(sharedKeyA));
        System.out.println("Bob   端共享密钥: " + bytesToHex(sharedKeyB));
        System.out.println("密钥一致? " + MessageDigest.isEqual(sharedKeyA, sharedKeyB));
    }

    private static String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder(bytes.length * 2);
        for (byte b : bytes)
            sb.append(String.format("%02x", b));
        return sb.toString();
    }
}